import numpy as np
import cv2 as cv
from matplotlib import pyplot as plt


# 通过叉积判断点在向量哪一侧(正左负右)
def cross_product(start, end, point):
    res = (end[0] - start[0]) * (point[1] - start[1]) - (end[1] - start[1]) * (point[0] - start[0])
    return res


# 图像处理的主函数
def draw_circles(ex, img):
    contours, hierarchy = cv.findContours(ex, cv.RETR_TREE, cv.CHAIN_APPROX_SIMPLE)
    # 筛选轮廓并对轮廓面积排序,第一大的为粗黑带,第二大的为细黑带
    contours = [cnt for cnt in contours if 200 < cv.contourArea(cnt) < 2000]
    area = []
    for cnt in contours:
        area.append(cv.contourArea(cnt))
    arg_area = np.argsort(area)
    thick_band = contours[arg_area[-1]]
    rect1 = cv.minAreaRect(thick_band)
    thick_center = rect1[0]
    box1 = cv.boxPoints(rect1)
    box1 = np.int64(box1)
    cv.drawContours(img, [box1], 0, (0, 0, 255), 2)
    slim_band = contours[arg_area[-2]]
    rect2 = cv.minAreaRect(slim_band)
    slim_center = rect2[0]
    box2 = cv.boxPoints(rect2)
    box2 = np.int64(box2)
    cv.drawContours(img, [box2], 0, (0, 0, 255), 2)
    # 用霍夫变换找圆
    circles = cv.HoughCircles(ex, cv.HOUGH_GRADIENT, 1, 10, param1=50, param2=9, minRadius=4, maxRadius=10)
    circles = np.uint16(np.around(circles))
    circles = circles[0, :]
    assert circles.shape[0] == 6, '找到的点数量错误!'
    # 用圆到粗黑带中心点的欧几里得距离对圆进行排序&两两分组成一行
    circles = sorted(circles, key=lambda x: (np.linalg.norm(x[:2] - thick_center)))
    circles = np.array(circles)
    circles = circles.reshape(3, 2, 3)
    # 以粗黑带中心点为起点,细黑带中心点为终点确定向量,通过与圆心的叉积确定点在向量哪一侧(正左负右)来行内排序
    sorted_circles = []
    for circles_row in circles:
        circles_row = sorted(circles_row, key=lambda x: (cross_product(thick_center, slim_center, x[:2])))
        sorted_circles.append(circles_row)
    sorted_circles = np.array(sorted_circles).reshape([6, 3])
    # 判断空心还是实心
    j = 0
    for i in sorted_circles:
        # 查看圆心周围像素点值,一旦有255便判定为空心
        a = np.int8(-2)
        while 1:
            if ex[i[1] + a, i[0] + a] == 255:
                cv.circle(img, (i[0], i[1]), i[2], (0, 255, 0), 2)
                cv.putText(img, '0,%d' % j, (i[0], i[1] + 20), cv.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 1)
                j = j + 1
                break
            a = a + 1
            if a == 3:
                cv.circle(img, (i[0], i[1]), i[2], (0, 0, 255), 2)
                cv.putText(img, '1,%d' % j, (i[0], i[1] + 20), cv.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 1)
                j = j + 1
                break


img1 = cv.imread('ex1.jpg')
ex1 = cv.cvtColor(img1, cv.COLOR_BGR2GRAY)
img2 = cv.imread('ex2.jpg')
ex2 = cv.cvtColor(img2, cv.COLOR_BGR2GRAY)
# 进行高斯模糊 核值小时会导致实心变空心
ex1 = cv.GaussianBlur(ex1, (9, 9), 0)
ex2 = cv.GaussianBlur(ex2, (9, 9), 0)
# 自适应阈值(均值) 使用高斯算法时也会导致实心变空心
ex1 = cv.adaptiveThreshold(ex1, 255, cv.ADAPTIVE_THRESH_MEAN_C, cv.THRESH_BINARY, 11, 15)
ex2 = cv.adaptiveThreshold(ex2, 255, cv.ADAPTIVE_THRESH_MEAN_C, cv.THRESH_BINARY, 11, 15)

# 调用主函数进行图像处理
draw_circles(ex1, img1)
draw_circles(ex2, img2)

# 改成rgb格式
b, g, r = cv.split(img1)
img1 = cv.merge([r, g, b])
b, g, r = cv.split(img2)
img2 = cv.merge([r, g, b])

plt.subplot(1, 2, 1), plt.imshow(img1), plt.xticks([]), plt.yticks([])
plt.subplot(1, 2, 2), plt.imshow(img2), plt.xticks([]), plt.yticks([])

plt.show()
